import {
  CitationDelta,
  MessageDelta,
  MessageStart,
  PacketType,
  StreamingCitation,
} from "./streamingModels";
import { Packet } from "@/app/chat/services/streamingModels";

export function isToolPacket(
  packet: Packet,
  includeSectionEnd: boolean = true
) {
  let toolPacketTypes = [
    PacketType.SEARCH_TOOL_START,
    PacketType.SEARCH_TOOL_QUERIES_DELTA,
    PacketType.SEARCH_TOOL_DOCUMENTS_DELTA,
    PacketType.PYTHON_TOOL_START,
    PacketType.PYTHON_TOOL_DELTA,
    PacketType.CUSTOM_TOOL_START,
    PacketType.CUSTOM_TOOL_DELTA,
    PacketType.REASONING_START,
    PacketType.REASONING_DELTA,
    PacketType.FETCH_TOOL_START,
    PacketType.FETCH_TOOL_URLS,
    PacketType.FETCH_TOOL_DOCUMENTS,
  ];
  if (includeSectionEnd) {
    toolPacketTypes.push(PacketType.SECTION_END);
  }
  return toolPacketTypes.includes(packet.obj.type as PacketType);
}

export function isDisplayPacket(packet: Packet) {
  return (
    packet.obj.type === PacketType.MESSAGE_START ||
    packet.obj.type === PacketType.IMAGE_GENERATION_TOOL_START ||
    packet.obj.type === PacketType.PYTHON_TOOL_START
  );
}

export function isSearchToolPacket(packet: Packet): boolean {
  return (
    packet.obj.type === PacketType.SEARCH_TOOL_START ||
    packet.obj.type === PacketType.SEARCH_TOOL_QUERIES_DELTA ||
    packet.obj.type === PacketType.SEARCH_TOOL_DOCUMENTS_DELTA
  );
}

export function isStreamingComplete(packets: Packet[]) {
  return packets.some((packet) => packet.obj.type === PacketType.STOP);
}

export function isFinalAnswerComing(packets: Packet[]) {
  return packets.some(
    (packet) =>
      packet.obj.type === PacketType.MESSAGE_START ||
      packet.obj.type === PacketType.IMAGE_GENERATION_TOOL_START ||
      packet.obj.type === PacketType.PYTHON_TOOL_START
  );
}

export function isFinalAnswerComplete(packets: Packet[]) {
  // Find the first MESSAGE_START packet and get its index
  const messageStartPacket = packets.find(
    (packet) =>
      packet.obj.type === PacketType.MESSAGE_START ||
      packet.obj.type === PacketType.IMAGE_GENERATION_TOOL_START ||
      packet.obj.type === PacketType.PYTHON_TOOL_START
  );

  if (!messageStartPacket) {
    return false;
  }

  // Check if there's a corresponding SECTION_END with the same turn_index
  return packets.some(
    (packet) =>
      packet.obj.type === PacketType.SECTION_END &&
      packet.turn_index === messageStartPacket.turn_index
  );
}

export function groupPacketsByTurnIndex(
  packets: Packet[]
): { turn_index: number; packets: Packet[] }[] {
  /*
  Group packets by turn_index. Ordered from lowest turn_index to highest turn_index.
  */
  const groups = packets.reduce((acc: Map<number, Packet[]>, packet) => {
    const turn_index = packet.turn_index;
    if (!acc.has(turn_index)) {
      acc.set(turn_index, []);
    }
    acc.get(turn_index)!.push(packet);
    return acc;
  }, new Map());

  // Convert to array and sort by turn_index (lowest to highest)
  return Array.from(groups.entries())
    .map(([turn_index, packets]) => ({
      turn_index,
      packets,
    }))
    .sort((a, b) => a.turn_index - b.turn_index);
}

export function getTextContent(packets: Packet[]) {
  return packets
    .map((packet) => {
      if (
        packet.obj.type === PacketType.MESSAGE_START ||
        packet.obj.type === PacketType.MESSAGE_DELTA
      ) {
        return (packet.obj as MessageStart | MessageDelta).content || "";
      }
      return "";
    })
    .join("");
}

export function getCitations(packets: Packet[]): StreamingCitation[] {
  const citations: StreamingCitation[] = [];

  packets.forEach((packet) => {
    if (packet.obj.type === PacketType.CITATION_DELTA) {
      const citationDelta = packet.obj as CitationDelta;
      citations.push(...(citationDelta.citations || []));
    }
  });

  return citations;
}
